import "dotenv/config";
import readline from "readline";
import { ChatOpenAI } from "@langchain/openai";
import { ChatMessageHistory } from "langchain/memory";
import { getBufferString } from "@langchain/core/messages";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import {
  RunnableSequence,
  RunnablePassthrough,
} from "@langchain/core/runnables";
import { StringOutputParser } from "@langchain/core/output_parsers";

const chatModel = new ChatOpenAI({
  apiKey: process.env.OPENAI_API_KEY,
  configuration: {
    baseURL: process.env.OPENAI_API_BASE,
  },
});

const chatPrompt = ChatPromptTemplate.fromMessages([
  [
    "system",
    `You are a helpful assistant. Answer all questions to the best of your ability.

  Here is the chat history summary:
  {history_summary}
  `,
  ],
  ["human", "{input}"],
]);

const summaryPrompt = ChatPromptTemplate.fromTemplate(`
  Progressively summarize the lines of conversation provided, adding onto the previous summary returning a new summary
  
  Current summary:
  {summary}
  
  New lines of conversation:
  {new_lines}
  
  New summary:
  `);

const summaryChain = RunnableSequence.from([
  summaryPrompt,
  chatModel,
  new StringOutputParser(),
]);

let summary = "";

const history = new ChatMessageHistory();

const chatChain = RunnableSequence.from([
  {
    input: new RunnablePassthrough({
      func: async (input) => {
        await history.addUserMessage(input);
        return input;
      },
    }),
  },
  RunnablePassthrough.assign({
    history_summary: () => summary,
  }),
  chatPrompt,
  chatModel,
  new StringOutputParser(),
  new RunnablePassthrough({
    func: async (output) => {
      // console.log("~~~~", input);
      await history.addAIMessage(output);
      const messages = await history.getMessages();
      const new_lines = getBufferString(messages);
      const newSummary = await summaryChain.invoke({
        summary,
        new_lines,
      });
      // history.clear();
      summary = newSummary;
      return output;
    },
  }),
]);

// CLI loop
const rl = readline.createInterface({
  input: process.stdin,
  output: process.stdout,
});

console.log("💬 开始对话吧！（输入 exit 可退出）");

const loop = () => {
  rl.question("> 你：", async (input) => {
    if (input.trim().toLowerCase() === "exit") {
      rl.close();
      return;
    }

    try {
      const res = await chatChain.invoke({ input });
      console.log("🤖 AI：", res);
      console.log("🧠 当前摘要：", summary);
    } catch (err) {
      console.error("❌ 错误：", err);
    }

    loop();
  });
};

loop();
